import pymssql
import sqlite3


class DbHelper:
    """
        数据处理类
    """
    cwd = ''
    # 进度表名
    PROGRESS_TB_NAME = 'tb_progress'

    def __init__(self, cwd):
        self.cwd = cwd

    def query_countries(self):
        """国家的编号，名字，英文名的列表"""
        countries = []
        conn = pymssql.connect(server='data.waimaobao.org',
                               port='9099',
                               user='sa',
                               password='HuL2277-_-!',
                               database='lancer_wbiso')
        cursor = conn.cursor()
        try:
            cursor.execute(
                "SELECT cid, country_name, country_name_en, dataformat"
                " FROM l_country_data WHERE dataformat <> 4 ORDER BY cid ASC")
            row = cursor.fetchone()
            while row:
                countries.append({"id": row[0], "name": row[1], "name_en": row[2], "data_format": row[3]})
                row = cursor.fetchone()
        except Exception as e:
            print(e)
        finally:
            cursor.close()
            conn.close()
        return countries

    def query_data(self, cid, name, data_type, index):
        """
        查询指定国家相关表的数据
        :param cid:     国家编号
        :param name:    国家名称
        :param data_type:    数据类型 (1- 进口 2- 出口)
        :param index:   开始索引
        :return: 数据的编号，公司名称
        """
        data_list = []
        conn = pymssql.connect(server='data.waimaobao.org',
                               port='9099',
                               user='sa',
                               password='HuL2277-_-!',
                               database='lancer_%s' % (self.get_db_name(name, cid)))
        cursor = conn.cursor()
        try:
            name_field = "importer" if data_type == 1 else "supplier"
            cursor.execute(
                "SELECT TOP 1000 id, %s, website, email1, email2, email3 FROM"
                " %s WHERE id > %d AND %s <> '' AND %s IS NOT NULL ORDER BY id ASC" %
                (name_field, "company_list_all" if data_type == 1 else "company_list_supplier", index, name_field,
                 name_field))
            row = cursor.fetchone()
            while row:
                data_list.append({"id": row[0], "name": row[1], "website": row[2], "email": [row[3], row[4], row[5]]})
                row = cursor.fetchone()
        except Exception as e:
            print(e)
        finally:
            cursor.close()
            conn.close()
        return {"data_list": data_list, "cid": cid, "name": name, "data_type": index}

    def get_db_name(self, name, cid):
        """获取数据库表的名称"""
        if cid == 6:
            name = 'KOREA'
        elif cid == 10:
            name = 'UK'
        elif cid == 11:
            name = 'UKRAIN'
        elif cid == 16:
            name = 'COSTA_RICA'
        elif cid == 24:
            name = 'EL' + name
        elif cid == 31:
            name = "USA_EX"
        return name

    def query_data_count(self, cid, name, data_type):
        """查询数据的总记录数"""
        conn = pymssql.connect(server='data.waimaobao.org',
                               port='9099',
                               user='sa',
                               password='HuL2277-_-!',
                               database='lancer_%s' % (self.get_db_name(name, cid)))
        cursor = conn.cursor()
        cursor.execute("SELECT max(id) FROM %s" % ("company_list_all" if data_type == 1 else "company_list_supplier"))
        row = cursor.fetchone()
        if row:
            return row[0]
        return 0

    def query_data_progress(self, cid, data_type):
        """查询数据搜索的进度"""
        progress = 0
        # ms_sql_server
        # conn = pymssql.connect(server='data.waimaobao.org',
        #                        port='9099',
        #                        user='sa',
        #                        password='HuL2277-_-!',
        #                        database='lancer_wbiso')
        # sqlite 3
        conn = self.get_table_conn(self.PROGRESS_TB_NAME)
        cursor = conn.cursor()
        try:
            # 查询进度
            cursor.execute("SELECT progress FROM %s WHERE cid = ? AND data_type = ?"
                           % self.PROGRESS_TB_NAME, (cid, data_type))
            row = cursor.fetchone()
            if row:
                progress = row[0]
        except Exception as e:
            print(e)
        finally:
            cursor.close()
            conn.close()
        return progress

    def update_data_progress(self, cid, data_type, progress):
        """更新数据搜索的进度"""
        conn = self.get_table_conn(self.PROGRESS_TB_NAME)
        try:
            conn.execute("REPLACE INTO %s VALUES (?, ?, ?)" % self.PROGRESS_TB_NAME, (cid, progress, data_type))
            conn.commit()
        except Exception as e:
            print(e)
        finally:
            conn.close()

    def update_data(self, country, data_type, result):
        conn = pymssql.connect(server='data.waimaobao.org',
                               port='9099',
                               user='sa',
                               password='HuL2277-_-!',
                               database='lancer_%s' % (self.get_db_name(country["name_en"], country["id"])))
        cursor = conn.cursor()
        try:
            sql = "UPDATE %s SET website = '%s', email1 = '%s', email2 = '%s', email3 = '%s' WHERE id = %d" \
                  % ("company_list_all" if data_type == 1 else "company_list_supplier",
                     result["site"], result["emails"][0], result["emails"][1],
                     result["emails"][2], result["id"])
            cursor.execute(sql)
            conn.commit()
        except Exception as e:
            print(e)
        finally:
            conn.close()

    def get_table_conn(self, tb_name):
        # 创建进度表
        conn = sqlite3.connect(self.cwd + '/local.db')
        try:
            conn.execute("CREATE TABLE IF NOT EXISTS %s("
                         "cid INTEGER,"
                         "progress INTEGER,"
                         "data_type INTEGER,"
                         "UNIQUE(cid, data_type))" % tb_name)
        except Exception as e:
            print(e)
        return conn
